import copy
import time
import traceback
import requests
from threading import Lock, Thread
from tqdm import tqdm
import os
import re
import time
import traceback

import requests
from tqdm import tqdm

debug = True
debug = False


class Downloader:
    def __init__(self, url, filename, headers=None, thread_num=8, strict=True):
        self.url = url
        self.filename = filename
        self.headers = headers
        if headers is None:
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.198 Safari/537.36',
                'Connection': 'keep-alive'
            }
        self.strict = strict
        if debug:
            print(f'[debug] 新建下载任务，文件名：{filename}，地址：{url}')

        self.cacheFileName = 'cache/' + self.validateTitle(self.filename) + '.cache'
        self.file = None
        self.file_write_lock = Lock()
        self.thread_num = thread_num
        self.one_block_size = 1024 * 64
        self.session = requests.Session()

        self.planned_size = 0
        self.planned_size_lock = Lock()

        self.downloading_size = 0
        self.downloading_size_lock = Lock()

        self.downloaded_size = 0

    def validateTitle(self, title):
        rstr = r"[/\\\:\*\?\"\<\>\|]"  # / \ : * ? " < > |
        new_title = re.sub(rstr, "_", title)
        return new_title

    def download(self, file_size):
        while self.planned_size < file_size:
            with self.planned_size_lock:
                start = self.planned_size
                end = min(self.planned_size + self.one_block_size, file_size)
                self.planned_size = end
            if start == end:
                break

            headers = copy.deepcopy(self.headers)
            headers['Range'] = f'bytes={start}-{end}'
            while True:
                try:
                    r = self.session.get(
                        self.url, stream=True, headers=headers, timeout=1)
                    with self.file_write_lock:
                        self.file.seek(start)
                        self.file.write(r.content)
                    break
                except requests.exceptions.ReadTimeout:
                    pass
                except requests.exceptions.ConnectionError:
                    time.sleep(0.1)
                    pass
                except requests.exceptions.ChunkedEncodingError:
                    pass
                except:
                    print(traceback.format_exc())
            with self.downloading_size_lock:
                self.downloading_size += end - start + 1

    def get_unit(self, file_size):
        unit_str = 'B'
        unit_int = 1
        if file_size >= 1024:
            file_size /= 1024
            unit_str = 'KB'
            unit_int = 1024
        if file_size >= 1024:
            file_size /= 1024
            unit_str = 'MB'
            unit_int = 1024 * 1024
        if file_size >= 1024:
            file_size /= 1024
            unit_str = 'GB'
            unit_int = 1024 * 1024 * 1024
        return unit_int, unit_str

    def monitor_thread(self, file_size):
        file_unit_int, file_unit_str = self.get_unit(file_size)
        with tqdm(total=file_size) as pbar:
            pbar.set_postfix(总大小='%.2f %s' % (
                file_size / file_unit_int, file_unit_str), 速度='0 KB/s')
            while self.downloaded_size < file_size:
                time.sleep(1)
                with self.downloading_size_lock:
                    self.downloaded_size += self.downloading_size
                    pbar.update(self.downloading_size)
                    downloading_unit_int, downloading_unit_str = self.get_unit(
                        self.downloading_size)
                    pbar.set_postfix(总大小='%.2f %s' % (
                        file_size / file_unit_int, file_unit_str), 速度='%.2f %s/s' % (self.downloading_size / downloading_unit_int, downloading_unit_str))
                    self.downloading_size = 0

    def download_payload(self):
        r2 = requests.get(self.url, stream=True, headers=self.headers)

        size = int(r2.headers['Content-Length'])
        if debug:
            print(f'[debug] 文件大小：{size}')

        time_beg = time.time()
        for i in range(self.thread_num):
            t1 = Thread(target=self.download, args=(size, ))
            t1.setDaemon(True)
            t1.start()

        self.monitor_thread(size)

        time_end = time.time()
        self.file.close()
        spendTime = time_end - time_beg

        return spendTime, size

    def main(self):
        if os.path.exists(self.filename):
            if self.strict:
                print(f'[{self.filename}]已存在，即将重新下载！')
                os.remove(self.filename)
            else:
                print(f'[{self.filename}]已存在，即将跳过下载！')
                return True
        self.file = open(self.cacheFileName, 'wb')
        try:
            spendTime, size = self.download_payload()
            speed_unit_int, speed_unit_str = self.get_unit(size / spendTime)
            os.rename(self.cacheFileName, self.filename)
            print(f'{self.filename}下载完成！用时 %.2f s 平均速度：%.2f %s/s' %
                  (spendTime, size / spendTime / speed_unit_int, speed_unit_str))
        except:
            print(f'{self.filename}下载出错！')
            print(traceback.format_exc())
            return False
        return True


if __name__ == '__main__':
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.198 Safari/537.36',
        'Connection': 'keep-alive'
    }
    filename = '仿佛打开异世界的大门.mp4'
    url = 'http://pan.zwtistudy.cn/?explorer/share/file&hash=94e1o8htDBcQTkcSz2PNJAWZcj9GdpQqnIfXmqR8vXbQNB4Eh2gSL7VL68N9Oz2T9Q'
    downloader = Downloader(url, filename=filename, headers=headers, thread_num=8, strict=True)
    downloader.main()
